{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Known error: This notebook requires [yt](https://yt-project.org/) to visualize the results. Yt needs to be updated to work properly first. Updates are currently being made to yt's frontends to make this PyNE integration work."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Let's Explore PyNE Meshes!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyne.xs.channels import sigma_t\n",
    "from pyne.material import Material, from_atom_frac\n",
    "from pyne.mesh import Mesh, NativeMeshTag, MetadataTag, ComputedTag\n",
    "from yt.config import ytcfg; ytcfg[\"yt\",\"suppressStreamLogging\"] = \"True\"\n",
    "from yt.mods import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's create a 2D structured mesh, whose grid points in x & y are [Cantor dust](http://en.wikipedia.org/wiki/Cantor_set#Cantor_dust).  Most of the mesh will have water as a coolant.  Some volume elements will have fuel, though."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cantor(n):\n",
    "    return [0.] + cant(0., 1., n) + [1.]\n",
    "\n",
    "def cant(x, y, n):\n",
    "    if n == 0:\n",
    "        return []\n",
    "    new_pts = [2.*x/3. + y/3., x/3. + 2.*y/3.]\n",
    "    return cant(x, new_pts[0], n-1) + new_pts + cant(new_pts[1], y, n-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "c5 = cantor(5)\n",
    "coords = [c5, c5, [0.0, 1.0]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We access volume elements (`ve`) on the mesh through a volume element index (`idx`).  This index is a unique integer `i` that gives a sort ordering to an otherwise unorder mesh.  The volume element index is defined on the range from 0 (inclusive) to the number of volumes in the mesh (exclusive).  You may access a volume element's material through the `mats` attribute and indexing with `i`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = Mesh(structured_coords=coords, structured=True)\n",
    "fuel = from_atom_frac({'U235': 0.045, 'U238': 0.955, 'O16': 2.0}, density=10.7)\n",
    "cool = from_atom_frac({'H1': 2.0, 'O16': 1.0}, density=1.0)\n",
    "for i in range(len(m)):\n",
    "    m.mats[i] = cool\n",
    "m.mats[len(m)/2] = fuel\n",
    "m.mats[len(m)/4] = fuel\n",
    "for i, c in enumerate(c5[:-1]):\n",
    "    m.mats[i*len(c5)] = fuel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tags"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Tag`s - sometimes known as fields - are generic way of storing data on a mesh.  Tags can be accessed as attributes on the mesh class itself.  For example, the density tag may be grabbed via:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MaterialPropertyTag(name='density', doc='the density [g/cc]')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.density "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get the value of the density, you have to provide the location by its index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.density[42]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also slice, mask, or fancy index tags:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 10.7   1.    1.    1.    1.    1.    1.    1.    1.    1.    1.    1.\n",
      "   1.    1.    1.    1.   10.7   1.    1.    1.    1.    1.    1.    1.\n",
      "   1.    1.    1.    1.    1.    1.    1.    1.   10.7   1.    1.    1.\n",
      "   1.    1.    1.    1. ]\n"
     ]
    }
   ],
   "source": [
    "print(m.density[::100])  # slice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7\n",
      "  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7\n",
      "  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7\n",
      "  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7\n",
      "  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7  10.7\n",
      "  10.7  10.7  10.7  10.7]\n"
     ]
    }
   ],
   "source": [
    "print(m.density[m.density[:] >= 10])  # mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[  1.   10.7   1.    1. ]\n"
     ]
    }
   ],
   "source": [
    "print(m.density[[10, 0, 11, 100]])  # fancy index is fancy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get a list of all current tag names, uses the `tags` dictionary on the mesh:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['sub_act',\n",
       " 'sub_lan',\n",
       " 'number_density',\n",
       " 'density',\n",
       " 'comp',\n",
       " 'expand_elements',\n",
       " 'to_atom_frac',\n",
       " 'molecular_weight',\n",
       " 'mass_density',\n",
       " 'atoms_per_mol',\n",
       " 'sub_tru',\n",
       " 'mass',\n",
       " 'attrs',\n",
       " 'mult_by_mass',\n",
       " 'sub_ma',\n",
       " 'idx',\n",
       " 'sub_fp']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.tags.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analysis & Visulaization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PyNE meshes are supported by the [yt project](http://yt-project.org/).  Use yt's plotting infrastructure to display tags on your mesh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"\"><br>"
      ],
      "text/plain": [
       "<yt.visualization.plot_window.AxisAlignedSlicePlot at 0x6155a10>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pf = PyneMoabHex8StaticOutput(m)\n",
    "s = SlicePlot(pf, 'z', 'density', origin='native')\n",
    "s.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computed Tags"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Computed tags - also known as 'derived fields' - are a way of having a lazily evaluated 'virtual tag.'  A computed tag is defined with a function, lambda, or any other callable object.  The function must have the following interface:\n",
    "\n",
    "```python\n",
    "def f(mesh, i):\n",
    "    \"\"\"mesh is a pyne.mesh.Mesh() object and i is the volume element index\n",
    "    to compute.\n",
    "    \"\"\"\n",
    "    # ... do some work ...\n",
    "    return anything_you_want\n",
    "```\n",
    "\n",
    "Here is a somewhat silly example which squares the density.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.density2 = ComputedTag(lambda mesh, i: mesh.density[i]**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the bounds on the color bar have changed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"\"><br>"
      ],
      "text/plain": [
       "<yt.visualization.plot_window.AxisAlignedSlicePlot at 0x6f31550>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pf = PyneMoabHex8StaticOutput(m)\n",
    "s = SlicePlot(pf, 'z', 'density2', origin='native')\n",
    "s.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a more serious example which uses PyNE's cross section tools to compute the one-group total cross section $\\sigma_t$ everywhere on the mesh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.sigma_t = ComputedTag(lambda mesh, i: sigma_t(mesh.mats[i], group_struct=[10.0, 1e-6], phi_g=[1.0])[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"\"><br>"
      ],
      "text/plain": [
       "<yt.visualization.plot_window.AxisAlignedSlicePlot at 0x70bee50>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pf = PyneMoabHex8StaticOutput(m)\n",
    "s = SlicePlot(pf, 'z', 'sigma_t', origin='native')\n",
    "s.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Yes, that was only one line of code."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
